# -*- coding: utf-8 -*-
"""
@Time    : 2024/7/11 15:42 
@Author  : ZhangShenao 
@File    : 1.模型IO.py 
@Desc    : Model IO
"""

import dotenv
from langchain_community.chat_models import ChatZhipuAI
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate

dotenv.load_dotenv()

# 创建原始模板
template = '''你是一位专业的鲜花店文案撰写员\n
对于售价为 {price} 元的 {flower_name}，请帮我生成一个简短且吸引人的宣传文案。
'''
# 根据原始模板,创建LangChain Prompt模板
prompt = ChatPromptTemplate.from_template(template)
# 创建模型实例
llm = ChatZhipuAI(model_name="glm-4-air")
# 创建输出解析器
parser = StrOutputParser()
# 构造Chain
chain = prompt | llm | parser

# 定义输入参数
prices = ['50', '30', '20']
flowers = ['玫瑰花', '百合花', '康乃馨']

# 调用LLM,获取输出结果
for price, flower in zip(prices, flowers):
    print(f'{flower}的文案: {chain.invoke({"price": price, "flower_name": flower})}')
